package com.paradogs.framework.core.annoations.rpc;

import com.paradogs.framework.core.events.RPCRespEvent;
import com.paradogs.framework.core.msg.PRMsg;
import com.paradogs.framework.core.msg.ByteData;
import com.paradogs.framework.core.msg.PRMsgContextHolder;
import com.paradogs.framework.core.utils.ProtoUtils;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: yumi
 * @date: 2023/8/11  13:45
 * @Description: RPC 响应处理
 */
@Slf4j
@Component
public class PRRPCResponseResolver {

    /**
     * 等待响应的消息 map
     */
    private volatile Map<Long, RespData> waitRespMap = new ConcurrentHashMap<>();


    /**
     * 监听 rpc 响应
     * @param event
     */
    @EventListener(RPCRespEvent.class)
    private void onRpcMsgResponse(RPCRespEvent event) {

        Long key = event.getMsg().getMsgHeader().getReplyMsgId();
        ByteData byteData = event.getMsg();

        PRRPCHolder.getMappingMap();
        try {
            // 执行回调方法
            RespData respInfo = waitRespMap.get(key);
            if (respInfo == null) {
                log.warn("Missing callback listener, maybe timeout");
                return;
            }
            if (respInfo.getPromise() == null) {
                log.warn("Promise == null");
                return;
            }

            // 解码消息体数据
            Object data = ProtoUtils.decode(byteData.getDataBytes(), respInfo.getReturnClass());

            // 设置上下文
            PRMsgContextHolder.setContext(respInfo.getMsgContextHolder());

            // 标记任务完成，执行回调
            boolean isSuccess = respInfo.getPromise().complete(data);
            if (!isSuccess) {log.error("Callback promise invoke is fail: {}", respInfo);}

        } catch (Exception e) {
            log.error("{}", e.getMessage(), e);
        } finally {
            // 处理完毕移除响应监听
            waitRespMap.remove(key);
            log.debug("Release callback listener: {}", key);
        }
    }

    /**
     * 监听一次
     * @param route
     * @param cgMsgId
     * @param promise
     */
    public <T> void on4Once(String route, Long cgMsgId, Class<?> returnClass, CompletableFuture<T> promise) {
        this.waitRespMap.put(cgMsgId, new RespData(returnClass, promise, PRMsgContextHolder.getSnapshot()));
    }

    @Data
    @AllArgsConstructor
    @NoArgsConstructor
    public static class RespData<T> {

        /**
         * 返回消息类型
         */
        private Class<? extends PRMsg> returnClass;

        /**
         * 回调方法
         */
        private CompletableFuture<T> promise;

        /**
         * 开始监听时的上下文状态
         */
        private PRMsgContextHolder.Snapshot msgContextHolder;
    }

}
